WIP produce TRX grouped by sl len #32
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces functionality to generate TRX tractography outputs grouped by streamline length bins, leveraging the existing GPU-based seeding/propagation pipeline.
Changes:
- Added
generate_trx_grouped_by_lenincu_tractography.pyto build multiple TRX files, one per length bin, while streaming seeds in chunks. - Extended
SeedBatchPropagatorincu_propagate_seeds.pywith binning utilities (gen_bin_indicesandas_array_sequence_group) to group generated streamlines by length across GPUs.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| cuslines/cuda_python/cu_tractography.py | Adds a TRX-generation path that partitions streamlines into multiple TRX files by length bins, using the new binning helpers. |
| cuslines/cuda_python/cu_propagate_seeds.py | Adds helper methods to compute per-bin streamline indices and expose grouped ArraySequence generators for use in binned TRX output. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| for bin_start in bin_starts: | ||
| max_steps = (bin_start + bin_len) / self.step_size | ||
| trx_files[bin_start] = TrxFile( | ||
| reference=ref_img, | ||
| nb_streamlines=n_sls_guess, | ||
| nb_vertices=n_sls_guess * max_steps, | ||
| ) |
There was a problem hiding this comment.
max_steps = (bin_start + bin_len) / self.step_size will generally be a float, so nb_vertices=n_sls_guess * max_steps passes a non-integer value where TrxFile expects a vertex count (an integer) and may use it to size numpy arrays/memmaps. To avoid type errors or under-allocation, consider computing an integer upper bound for the number of steps per streamline (e.g., using math.ceil on the division) and passing that integer to nb_vertices.
| def gen_bin_indices(self, bin_starts, bin_len): | ||
| bin_edges = np.append(bin_starts, bin_starts[-1] + bin_len) | ||
| bin_indices = {k: [] for k in range(len(bin_starts))} | ||
|
|
||
| for ii in range(self.ngpus): | ||
| scaled_lens = self.sline_lens[ii] * self.gpu_tracker.step_size | ||
| assignments = np.digitize(scaled_lens, bin_edges) - 1 | ||
|
|
||
| for k in range(len(bin_starts)): | ||
| jj_indices = np.where(assignments == k)[0] | ||
| if jj_indices.size > 0: | ||
| bin_indices[k].append((ii, jj_indices)) | ||
|
|
||
| return bin_indices | ||
|
|
||
| def as_array_sequence_group(self, bin_indices, bin_start): | ||
| relevant_blocks = bin_indices[bin_start] | ||
|
|
There was a problem hiding this comment.
gen_bin_indices builds bin_indices as a dict keyed by integer indices 0..len(bin_starts)-1, but as_array_sequence_group indexes bin_indices using the bin_start argument directly; when callers pass the physical bin start values (e.g., 20, 30, ...), this will not match the integer keys and will lead to KeyError or empty groups. The keys should be made consistent (either indices everywhere or the actual bin_start values), for example by keying bin_indices by the bin start values inside gen_bin_indices and keeping as_array_sequence_group's interface as-is.
| for bin_start in bin_starts: | ||
| tractogram = Tractogram( | ||
| self.seed_propagator.as_array_sequence_group(bin_indices, bin_start), |
There was a problem hiding this comment.
Here bin_indices = self.seed_propagator.gen_bin_indices(bin_starts, bin_len) returns a dict keyed by integer bin indices (0..len(bin_starts)-1), but the subsequent loop uses the physical bin_start values (e.g., 20, 30, ...) as keys when calling as_array_sequence_group, which internally does bin_indices[bin_start]. This mismatch in key types means lookups into bin_indices will fail or return no streamlines; the binning logic needs to use consistent keys (either indices or bin start values) across gen_bin_indices, generate_trx_grouped_by_len, and as_array_sequence_group.
| for bin_start in bin_starts: | |
| tractogram = Tractogram( | |
| self.seed_propagator.as_array_sequence_group(bin_indices, bin_start), | |
| for bin_idx, bin_start in enumerate(bin_starts): | |
| tractogram = Tractogram( | |
| self.seed_propagator.as_array_sequence_group(bin_indices, bin_idx), |
1 participant
No description provided.